# sklearn中的KNN分类器
import numpy as np
from sklearn.neighbors import KNeighborsClassifier
import matplotlib.pyplot as plt

data = np.loadtxt("C:\\Users\ASUS\Desktop\机器学习\数据集\gaussian_data.csv", delimiter=',')
data_train = data[:, :2]
label_train = data[:, 2]
print('数据集大小：', len(data_train))

plt.figure()
plt.scatter(data_train[label_train == 0, 0], data_train[label_train == 0, 1], c=[(1, 0, 0, 1)], marker='o')
plt.scatter(data_train[label_train == 1, 0], data_train[label_train == 1, 1], c=[(0, 1, 0, 1)], marker='x')
plt.scatter(data_train[label_train == 2, 0], data_train[label_train == 2, 1], c=[(0, 0, 1, 1)], marker='.')

plt.xlabel('X axis')
plt.ylabel('Y axis')
plt.show()

min_x, max_x = np.min(data_train[:, 0]) - 1, np.max(data_train[:, 0]) + 1
min_y, max_y = np.min(data_train[:, 1]) - 1, np.max(data_train[:, 1]) + 1

data_test_x = np.arange(min_x, max_x, 0.2)
data_test_y = np.arange(min_y, max_y, 0.2)
data_test = []
for i in data_test_x:
    for j in data_test_y:
        data_test.append(np.asarray([i, j]))
data_test = np.asarray(data_test)

ks = [1, 3, 5, 7, 9, 100]
for i, k in enumerate(ks):
    knn = KNeighborsClassifier(n_neighbors=k)
    knn.fit(data_train, label_train)
    label_test = knn.predict(data_test)

    plt.figure(figsize=(max_x - min_x + 1, max_y - min_y + 1))
    plt.scatter(data_train[label_train == 0, 0], data_train[label_train == 0, 1], c=[(1, 0, 0, 1)], marker='o')
    plt.scatter(data_train[label_train == 1, 0], data_train[label_train == 1, 1], c=[(0, 1, 0, 1)], marker='x')
    plt.scatter(data_train[label_train == 2, 0], data_train[label_train == 2, 1], c=[(0, 0, 1, 1)], marker='.')

    plt.scatter(data_test[label_test == 0, 0], data_test[label_test == 0, 1], c=[(1, 0, 0, 0.5)], marker='^')
    plt.scatter(data_test[label_test == 1, 0], data_test[label_test == 1, 1], c=[(0, 1, 0, 0.5)], marker='<')
    plt.scatter(data_test[label_test == 2, 0], data_test[label_test == 2, 1], c=[(0, 0, 1, 0.5)], marker='>')

    plt.title('k=' + str(k))
    plt.xlabel('X axis')
    plt.ylabel('Y axis')
    plt.show()

